package com.example.demo.ChatMemory; // 使用您提供的包名

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.util.Assert;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * 基于 Kryo 文件持久化的 ChatMemory。
 * 显式注册类并使用 writeClassAndObject/readClassAndObject。
 */
public class KryoFileChatMemory implements ChatMemory {

    private final Path storagePath;
    private static final String FILE_EXTENSION = ".kryo";
    private static final Kryo kryo = new Kryo();

    static {
        // --- Kryo 最佳实践配置 ---

        // 1. 设置实例化策略，以支持没有无参构造函数的类（关键修复点）
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());

        // 2. 强制要求类必须注册，这能提高性能并避免很多隐晦的错误
        kryo.setRegistrationRequired(false);

        // 3. 显式注册所有需要序列化的类
        kryo.register(ArrayList.class);
        kryo.register(HashMap.class); // Message 中的 properties 可能是 HashMap
        kryo.register(UserMessage.class);
        kryo.register(AssistantMessage.class);
        kryo.register(SystemMessage.class);
        // 如果你还使用了其他 Message 类型或它们包含了其他复杂对象，也需要在此注册
    }

    public KryoFileChatMemory(String storageDirectoryPath) {
        Assert.hasText(storageDirectoryPath, "存储目录路径不能为空");
        this.storagePath = Paths.get(storageDirectoryPath);
        if (!Files.exists(this.storagePath)) {
            try {
                Files.createDirectories(this.storagePath);
            } catch (IOException e) {
                throw new RuntimeException("无法创建存储目录: " + this.storagePath, e);
            }
        }
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        Assert.hasText(conversationId, "conversationId 不能为空");
        Assert.notNull(messages, "消息列表不能为空");

        List<Message> allMessages = get(conversationId);
        allMessages.addAll(messages);
        saveMessages(conversationId, allMessages);
    }

    @Override
    @SuppressWarnings("unchecked")
    public List<Message> get(String conversationId) {
        Assert.hasText(conversationId, "conversationId 不能为空");

        File conversationFile = getFileForConversation(conversationId);
        if (!conversationFile.exists() || conversationFile.length() == 0) {
            return new ArrayList<>();
        }

        try (Input input = new Input(new FileInputStream(conversationFile))) {
            // 使用 readClassAndObject
            return (List<Message>) kryo.readClassAndObject(input);
        } catch (Exception e) {
            // 此处如果再出错，通常是由于有未注册的类导致。
            // 请检查异常栈 e.getMessage()，看是否提示 "Class is not registered"。
            throw new RuntimeException("从文件加载会话失败: " + conversationFile.getAbsolutePath(), e);
        }
    }

    @Override
    public void clear(String conversationId) {
        File fileForConversation = getFileForConversation(conversationId);
        if (fileForConversation.exists()) {
            fileForConversation.delete();
        }
    }

    private void saveMessages(String conversationId, List<Message> messages) {
        File conversationFile = getFileForConversation(conversationId);
        try (Output output = new Output(new FileOutputStream(conversationFile))) {
            // 使用 writeClassAndObject，将类的元数据和对象数据一起写入
            kryo.writeClassAndObject(output, messages);
        } catch (Exception e) {
            // 此处如果出错，通常是由于有未注册的类导致。
            throw new RuntimeException("保存会话到文件失败: " + conversationFile.getAbsolutePath(), e);
        }
    }

    private File getFileForConversation(String conversationId) {
        String sanitizedId = conversationId.replaceAll("[^a-zA-Z0-9_\\-]", "");
        return this.storagePath.resolve(sanitizedId + FILE_EXTENSION).toFile();
    }
}